# -*- coding: utf-8 -*- 
# @Time : 2022/4/4 16:24 
# @Author : zzuxyj 
# @File : 15-train-GPU02.py
# -*- coding: utf-8 -*-
# @Time : 2022/4/4 16:06
# @Author : zzuxyj
# @File : 15-train-GPU01.py


"""
完整实现模型训练 在GPU上训练 常用写法
"""
import torch
import torchvision
import time

# 加载数据集
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

trainSet = torchvision.datasets.CIFAR10("../dataset/CIFAR10", train=True, transform=torchvision.transforms.ToTensor(),
                                        download=True)
testSet = torchvision.datasets.CIFAR10("../dataset/CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)

# 计算长度
print("训练集长度：{}".format(len(trainSet)))
print("测试集长度：{}".format(len(testSet)))

# 加载数据集合
trainSetLoad = DataLoader(trainSet, batch_size=64)
testSetLoad = DataLoader(testSet, batch_size=64)

# 创建网络模型  vgg16
vgg16Param = torch.load("./model/vgg16-pretrain.pth")
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(vgg16Param)
vgg16.classifier[6] = torch.nn.Linear(4096, 10)  # 10分类 修改网络模型

# 计算位置
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
vgg16.to(device)

# 损失函数和优化器
loss_fn = torch.nn.CrossEntropyLoss()
loss_fn.to(device)
# 学习率
learning_rate = 1e-2
optimizer = torch.optim.SGD(vgg16.parameters(), lr=learning_rate)
# 调整学习率
schduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.1)

# 训练网络的参数
# 训练次数
total_train_step = 0
# 训练轮次
train_epoch = 20
# 测试次数
test_train_step = 0

# tensorboard
writer = SummaryWriter("logs14")

# 训练开始
for epoch in range(train_epoch):
    print("-----------第{}轮次训练开始-------------".format(epoch + 1))
    # 训练步骤开始
    vgg16.train()
    start_time = time.time()
    # 单轮次训练开始
    for data in trainSetLoad:
        imgs, target = data
        imgs = imgs.to(device)
        target = target.to(device)
        output = vgg16(imgs)
        result_loss = loss_fn(output, target)

        # 优化网络模型
        optimizer.zero_grad()
        result_loss.backward()

        total_train_step += 1
        if (total_train_step % 100 == 0):
            end_time = time.time()
            print("训练时间：{}".format(end_time - start_time))
            print("训练次数：{},Loss：{}".format(total_train_step, result_loss.item()))
            writer.add_scalar("train_loss", result_loss.item(), total_train_step)

    # 优化lr
    schduler.step()

    # 测试步骤开始
    vgg16.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad(): # torch.no_grad() 梯度设置为空 ，然后进行测试
        for data in testSetLoad:
            imgs, target = data
            imgs = imgs.to(device)
            target = target.to(device)
            output = vgg16(imgs)
            test_loss = loss_fn(output, target)
            total_test_loss = total_test_loss + test_loss
            accuracy = (output.argmax(1) == target).sum()
            total_accuracy = total_accuracy + accuracy

    print("测试集上的Loss：{}".format(total_test_loss))
    print("测试集上的正确率：{}".format(total_accuracy / len(testSet)))
    writer.add_scalar("test_loss", total_test_loss, test_train_step)
    writer.add_scalar("test_accuracy", total_accuracy / len(testSet), test_train_step)
    test_train_step += 1

    if ((total_accuracy / len(testSet)) > 0.8):
        torch.save(vgg16.state_dict(), "./model/vgg16-model{}-CIFAR10.pth".format(epoch))
        print("模型保存")

# 关闭流
writer.close()

